package main

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/pem"
	"fmt"
	"math/big"
	"net"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"syscall"
	"time"

	"golang.org/x/term"
)

func checkOrGenerateCertificate(domain string, config Config) (tls.Certificate, error) {
	certFile := filepath.Join("cert", domain+".crt")
	keyFile := filepath.Join("cert", domain+".key")
	
	// 检查证书文件是否存在
	if _, err := os.Stat(certFile); err == nil {
		if _, err := os.Stat(keyFile); err == nil {
			fmt.Printf("使用现有证书: %s\n", certFile)
			return tls.LoadX509KeyPair(certFile, keyFile)
		}
	}
	
	fmt.Printf("为域名 %s 在内存中生成证书\n", domain)
	return generateCertificateInMemory(domain, config)
}

func generateCertificateInMemory(domain string, config Config) (tls.Certificate, error) {
	// 生成私钥
	privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		return tls.Certificate{}, fmt.Errorf("生成私钥失败: %v", err)
	}
	
	// 创建证书模板
	template := x509.Certificate{
		SerialNumber: big.NewInt(1),
		Subject: pkix.Name{
			Organization:  []string{"Proxy Server"},
			Country:       []string{"US"},
			Province:      []string{"CA"},
			Locality:      []string{"San Francisco"},
			StreetAddress: []string{""},
			PostalCode:    []string{""},
			CommonName:    domain,
		},
		NotBefore:    time.Now(),
		NotAfter:     time.Now().Add(365 * 24 * time.Hour),
		KeyUsage:     x509.KeyUsageKeyEncipherment | x509.KeyUsageDigitalSignature,
		ExtKeyUsage:  []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
		IPAddresses:  []net.IP{net.IPv4(127, 0, 0, 1)},
		DNSNames:     []string{domain},
	}
	
	// 从嵌入文件系统或指定路径加载CA证书和私钥
	var caCertPEM, caKeyPEM []byte
	
	// 优先使用命令行指定的CA文件
	if config.CAFile != "" && config.CAFile != "cert/ca.pem" {
		caCertPEM, err = os.ReadFile(config.CAFile)
		if err != nil {
			return tls.Certificate{}, fmt.Errorf("读取指定CA证书失败: %v", err)
		}
		
		caKeyPEM, err = os.ReadFile(config.CAKeyFile)
		if err != nil {
			return tls.Certificate{}, fmt.Errorf("读取指定CA私钥失败: %v", err)
		}
	} else {
		// 回退到嵌入文件系统
		caCertPEM, err = certFS.ReadFile("cert/ca.pem")
		if err != nil {
			return tls.Certificate{}, fmt.Errorf("读取嵌入的CA证书失败: %v", err)
		}
		
		caKeyPEM, err = certFS.ReadFile("cert/ca.key")
		if err != nil {
			return tls.Certificate{}, fmt.Errorf("读取嵌入的CA私钥失败: %v", err)
		}
	}
	
	caCertBlock, _ := pem.Decode(caCertPEM)
	caCert, err := x509.ParseCertificate(caCertBlock.Bytes)
	if err != nil {
		return tls.Certificate{}, fmt.Errorf("解析CA证书失败: %v", err)
	}
	
	caKeyBlock, _ := pem.Decode(caKeyPEM)
	caKeyInterface, err := x509.ParsePKCS8PrivateKey(caKeyBlock.Bytes)
	if err != nil {
		// 尝试PKCS1格式作为备选
		caKeyInterface, err = x509.ParsePKCS1PrivateKey(caKeyBlock.Bytes)
		if err != nil {
			return tls.Certificate{}, fmt.Errorf("解析CA私钥失败: %v", err)
		}
	}
	
	caKey, ok := caKeyInterface.(*rsa.PrivateKey)
	if !ok {
		return tls.Certificate{}, fmt.Errorf("CA私钥不是RSA格式")
	}
	
	// 使用CA签名生成证书
	certDER, err := x509.CreateCertificate(rand.Reader, &template, caCert, &privateKey.PublicKey, caKey)
	if err != nil {
		return tls.Certificate{}, fmt.Errorf("生成证书失败: %v", err)
	}
	
	// 将证书和私钥转换为PEM格式
	certPEM := pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: certDER})
	keyPEM := pem.EncodeToMemory(&pem.Block{Type: "RSA PRIVATE KEY", Bytes: x509.MarshalPKCS1PrivateKey(privateKey)})
	
	// 从内存中加载证书
	cert, err := tls.X509KeyPair(certPEM, keyPEM)
	if err != nil {
		return tls.Certificate{}, fmt.Errorf("加载内存证书失败: %v", err)
	}
	
	fmt.Printf("成功在内存中生成域名 %s 的证书\n", domain)
	return cert, nil
}

func checkAndInstallCARoot(config Config) error {
	var caCertPEM []byte
	var err error
	
	// 优先使用命令行指定的CA文件
	if config.CAFile != "" && config.CAFile != "cert/ca.pem" {
		caCertPEM, err = os.ReadFile(config.CAFile)
		if err != nil {
			return fmt.Errorf("读取指定CA证书失败: %v", err)
		}
	} else {
		// 优先从嵌入文件系统读取CA证书
		caCertPEM, err = certFS.ReadFile("cert/ca.pem")
		if err != nil {
			// 回退到本地文件系统
			caCertFile := filepath.Join("cert", "ca.pem")
			if _, err := os.Stat(caCertFile); os.IsNotExist(err) {
				return fmt.Errorf("CA证书文件不存在: %s", caCertFile)
			}
			
			caCertPEM, err = os.ReadFile(caCertFile)
			if err != nil {
				return fmt.Errorf("读取CA证书失败: %v", err)
			}
		}
	}
	
	caCertBlock, _ := pem.Decode(caCertPEM)
	caCert, err := x509.ParseCertificate(caCertBlock.Bytes)
	if err != nil {
		return fmt.Errorf("解析CA证书失败: %v", err)
	}
	
	// 检查CA证书是否已在系统中被信任
	if isCAInstalled(caCert) {
		fmt.Println("CA证书已在系统中被信任")
		return nil
	}
	
	fmt.Println("CA证书未被信任，需要安装到系统钥匙串...")
	return installCAToKeychainFromMemory(caCertPEM)
}

func isCAInstalled(caCert *x509.Certificate) bool {
	// 使用security命令检查证书是否已安装
	cmd := exec.Command("security", "find-certificate", "-c", caCert.Subject.CommonName, "/System/Library/Keychains/SystemRootCertificates.keychain")
	err := cmd.Run()
	if err == nil {
		return true
	}
	
	// 检查登录钥匙串
	cmd = exec.Command("security", "find-certificate", "-c", caCert.Subject.CommonName)
	err = cmd.Run()
	return err == nil
}

func installCAToKeychainFromMemory(caCertPEM []byte) error {
	fmt.Print("需要管理员权限安装CA证书，请输入密码: ")
	password, err := term.ReadPassword(int(syscall.Stdin))
	if err != nil {
		return fmt.Errorf("读取密码失败: %v", err)
	}
	fmt.Println()
	
	// 创建临时文件
	tmpFile, err := os.CreateTemp("", "ca-cert-*.pem")
	if err != nil {
		return fmt.Errorf("创建临时文件失败: %v", err)
	}
	defer os.Remove(tmpFile.Name())
	defer tmpFile.Close()
	
	// 写入证书内容
	if _, err := tmpFile.Write(caCertPEM); err != nil {
		return fmt.Errorf("写入临时文件失败: %v", err)
	}
	tmpFile.Close()
	
	// 安装证书到系统钥匙串
	cmd := exec.Command("sudo", "-S", "security", "add-trusted-cert", "-d", "-r", "trustRoot", "-k", "/Library/Keychains/System.keychain", tmpFile.Name())
	cmd.Stdin = strings.NewReader(string(password) + "\n")
	
	output, err := cmd.CombinedOutput()
	if err != nil {
		return fmt.Errorf("安装CA证书失败: %v, 输出: %s", err, string(output))
	}
	
	fmt.Printf("成功安装CA证书到系统钥匙串\n")
	return nil
}